#!/bin/bash

# 定义颜色
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
RED='\033[0;31m'
NC='\033[0m' # 无颜色

echo -e "${YELLOW}启动本地ONNX模型 - 医疗NLP工具箱...${NC}"

# 设置环境变量
export PYTHONPATH=$PYTHONPATH:$(pwd)
export NER_MODEL_PATH="/appslog/NER/Medical-NER-ONNX"
export XINFERENCE_SERVER_URL="http://10.128.10.186:9997"
export HF_HUB_DISABLE_SYMLINKS_WARNING=1
export HF_HUB_DISABLE_TELEMETRY=1
export TRANSFORMERS_CACHE="/tmp/transformers_cache"

# 创建缓存目录
mkdir -p /tmp/transformers_cache

# 检查ONNX模型路径
if [ -d "$NER_MODEL_PATH" ]; then
    echo -e "${GREEN}找到本地ONNX模型目录: $NER_MODEL_PATH${NC}"
    echo -e "${GREEN}目录内容:${NC}"
    ls -la $NER_MODEL_PATH
else
    echo -e "${RED}错误: 未找到本地ONNX模型目录 $NER_MODEL_PATH${NC}"
    echo -e "${RED}请确保路径正确或创建此目录并下载模型文件${NC}"
    exit 1
fi

# 检查ONNX模型文件
ONNX_MODEL_FILE="$NER_MODEL_PATH/onnx/model.onnx"
if [ -f "$ONNX_MODEL_FILE" ]; then
    echo -e "${GREEN}找到ONNX模型文件: $ONNX_MODEL_FILE${NC}"
else
    echo -e "${RED}错误: 未找到ONNX模型文件 $ONNX_MODEL_FILE${NC}"
    echo -e "${RED}请确保模型文件已正确下载到指定位置${NC}"
    exit 1
fi

# 检查分词器文件
TOKENIZER_FILES_FOUND=false

if [ -f "$NER_MODEL_PATH/tokenizer.json" ]; then
    echo -e "${GREEN}找到分词器配置文件: $NER_MODEL_PATH/tokenizer.json${NC}"
    TOKENIZER_FILES_FOUND=true
fi

if [ -f "$NER_MODEL_PATH/tokenizer_config.json" ]; then
    echo -e "${GREEN}找到分词器配置文件: $NER_MODEL_PATH/tokenizer_config.json${NC}"
    TOKENIZER_FILES_FOUND=true
fi

if [ -f "$NER_MODEL_PATH/spm.model" ]; then
    echo -e "${GREEN}找到SentencePiece模型文件: $NER_MODEL_PATH/spm.model${NC}"
    TOKENIZER_FILES_FOUND=true
fi

if [ -f "$NER_MODEL_PATH/config.json" ]; then
    echo -e "${GREEN}找到模型配置文件: $NER_MODEL_PATH/config.json${NC}"
else
    echo -e "${YELLOW}警告: 未找到模型配置文件 $NER_MODEL_PATH/config.json${NC}"
    echo -e "${YELLOW}将使用默认标签映射${NC}"
fi

if [ "$TOKENIZER_FILES_FOUND" = false ]; then
    echo -e "${RED}错误: 未在 $NER_MODEL_PATH 中找到必要的分词器文件${NC}"
    echo -e "${RED}请确保至少有tokenizer.json或tokenizer_config.json文件${NC}"
    exit 1
fi

# 检查数据库文件
DB_PATH="backend/db/snomed_bge_m3.db"
if [ -f "$DB_PATH" ]; then
    echo -e "${GREEN}找到数据库文件: $DB_PATH${NC}"
else
    echo -e "${YELLOW}警告: 未找到数据库文件 $DB_PATH, 术语标准化功能可能无法正常工作${NC}"
fi

# 安装依赖
echo -e "${GREEN}安装兼容版本的transformers和依赖...${NC}"
pip install transformers==4.28.0 tokenizers==0.13.3 onnxruntime numpy

echo -e "${GREEN}显示已安装的transformers版本:${NC}"
pip show transformers

echo -e "${GREEN}显示已安装的tokenizers版本:${NC}"
pip show tokenizers

# 启动后端服务
echo -e "${GREEN}启动后端服务...${NC}"
cd backend
python main.py &
BACKEND_PID=$!
cd ..

# 等待后端启动
echo "等待后端服务启动..."
sleep 5

# 检查后端是否启动成功
if kill -0 $BACKEND_PID 2>/dev/null; then
    echo -e "${GREEN}后端服务已启动，PID: $BACKEND_PID${NC}"
else
    echo -e "${RED}错误: 后端服务启动失败${NC}"
    echo -e "${YELLOW}请检查日志以获取更多信息${NC}"
    exit 1
fi

# 启动前端应用
echo -e "${GREEN}启动前端应用...${NC}"
cd frontend
npm start &
FRONTEND_PID=$!
cd ..

# 打印访问信息
echo -e "${GREEN}==========================${NC}"
echo -e "${GREEN}医疗NLP工具箱已启动${NC}"
echo -e "${GREEN}后端API: ${NC}http://localhost:8000"
echo -e "${GREEN}前端应用: ${NC}http://localhost:3000"
echo -e "${GREEN}使用本地ONNX模型: ${NC}${ONNX_MODEL_FILE}"
echo -e "${GREEN}使用xinference服务: ${NC}${XINFERENCE_SERVER_URL}"
echo -e "${YELLOW}按Ctrl+C停止服务${NC}"
echo -e "${GREEN}==========================${NC}"

# 捕获CTRL+C信号
trap "echo -e '${YELLOW}正在停止服务...${NC}'; kill $BACKEND_PID $FRONTEND_PID; echo '服务已停止'; exit 0" INT

# 保持脚本运行
wait 